#%%
import pickle
import numpy as np
import pandas as pd
import umap
import os
import os.path as osp
import matplotlib.pyplot as plt
from lilab.OpenLabCluster_train.a1_mirror_mutual_filt_clippredpkl import factory_label_mirror_start0

clippredpkl='/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/semiseq2seq_iter0/output/far_ns_with_s_recluster_k34/olc-2024-05-23-semiseq2seq-iter2-epoch5_pca12_farns_withs.clippredpkl'
clippreddata = pickle.load(open(clippredpkl, 'rb'))
df_clipNames = clippreddata['df_clipNames'].copy()
df_clipNames['index_raw'] = df_clipNames.index
df_clipNames['cluster_labels'] = clippreddata['cluster_labels']

# %%
_, fun_label_mirror = factory_label_mirror_start0(clippreddata['nK_mutual'], clippreddata['nK_mirror_half'])
df_clipNames_sort = df_clipNames.sort_values(by=['vnake', 'startFrame', 'isBlack'])
cluster_labels2 = df_clipNames_sort['cluster_labels'].values
cluster_labels3 = cluster_labels2.reshape(-1,2)[:,::-1].ravel()
df_clipNames_sort['cluster_labels_reverse'] = cluster_labels3
df_clipNames_sort['cluster_labels_mirror'] = fun_label_mirror(cluster_labels2-1) + 1
is_mirror = df_clipNames_sort['cluster_labels_mirror'] == df_clipNames_sort['cluster_labels_reverse']
m = is_mirror.mean()
assert m>0.7, 'mirror label is not good'
print(f'mirror label is good, percentage:{m:.2f}', )

#%%
df_clipNames_mirror = df_clipNames_sort[is_mirror]
index_raw = df_clipNames_mirror['index_raw'].values
embedding = clippreddata['embedding'][index_raw]
cluster_labels = clippreddata['cluster_labels'][index_raw]
clipNames = clippreddata['clipNames'][index_raw]
df_clipNames = clippreddata['df_clipNames'].loc[index_raw].reset_index(drop=True)

reducer = umap.UMAP(random_state=1000)
embedding_d2 = reducer.fit_transform(embedding)
clippreddata_new = clippreddata.copy()
clippreddata_new.update({'embedding_d2':embedding_d2, 
                         'embedding':embedding,
                         'cluster_labels':cluster_labels,
                         'df_clipNames':df_clipNames,
                         'clipNames':clipNames})

output_dir = osp.join(osp.dirname(clippredpkl), f'filt_perc{int(m*100)}')
os.makedirs(output_dir, exist_ok=True)
outdata_path = osp.join(output_dir, osp.basename(clippredpkl))
pickle.dump(clippreddata_new, open(outdata_path, 'wb'))

plt.figure(figsize=(12, 10))
plt.scatter(embedding_d2[::2, 0], embedding_d2[::2, 1], c=cluster_labels[::2], s=1, cmap='hsv') #
plt.colorbar()
plt.title(f'Repr. Proportion {len(cluster_labels)} 100% (all)')
plt.xticks([])
plt.yticks([])
plt.savefig(osp.join(output_dir, f'olc-2024-05-23-semiseq2seq_pca12_farns_withs.png'))
